package com.feidee.fd.sml.algorithm.ml

import com.feidee.fd.sml.algorithm.component.ml.regression.{DTRegressionComponent, DTRegressionParam}
import com.feidee.fd.sml.algorithm.util.{TestingDataGenerator, ToolClass}
import org.scalatest.FunSuite

/**
  * @Author songhaicheng
  * @Date 2019/3/26 15:17
  * @Description
  * @Reviewer
  */
class DTRegressionComponentSuite extends FunSuite {

  val dtr = new DTRegressionComponent
  val paramStr: String =
    """
      |{
      |	'maxBins': 3,
      |	'minInfoGain': 0.01
      |}
    """.stripMargin
  val param: DTRegressionParam = dtr.parseParam(new ToolClass().encrypt(paramStr))

  test("DecisionTreeRegression parameter") {
    assert(param.varianceCol == null && "variance".equals(param.impurity) && param.maxDepth == 5 &&
      param.maxBins == 3 && param.minInstancesPerNode == 1 && param.checkpointInterval == 10 &&
      param.maxMemoryInMB == 256 && param.seed == 123456L && param.minInfoGain == 0.01 && !param.cacheNodeIds)
  }

  test("DecisionTreeRegression prediction") {
    val data = TestingDataGenerator.sampleLinearRegressionData
    val res = dtr.train(param, data).transform(data)
    assert(res.schema.size == 3)
  }

}
